package com.practice.zuocy.base.class10;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * 并查集：
 * （1）有若干个样本a、b、c、d...类型假设是V
 * （2）在并查集中一开始认为每个样本都在单独的集合里
 * （3）用户可以在任何时候调用如下两个方法：
 *     boolean isSameSet(V x, V y)：查询样本x和样本y是否属于一个集合
 *     void union(V x, V y)：把x和y各自所在集合的所有样本合并成一个集合
 * （4）isSameSet和union方法的代价越低越好
 */
public class Code01_UnionFind {

	// 集合节点类型
	public static class Node<V> {
		V value;

		public Node(V v) {
			value = v;
		}
	}

	// 并查集结构
	public static class UnionSet<V> {

		// 集合中的节点Map：值为key的的节点是value节点。V -> 节点
		public HashMap<V, Node<V>> nodes;

		// 节点的父节点映射：key的父节点为value
		public HashMap<Node<V>, Node<V>> parents;

		// 只有一个节点是代表节点，才在size表中有记录
		// 代表节点key所在的集合有value个节点（集合的size信息）
		public HashMap<Node<V>, Integer> sizeMap;

		public UnionSet(List<V> values) {
			for (V cur : values) {
				Node<V> node = new Node<>(cur);
				// 初始时每个集合只有一个节点，就是代表节点（每个节点都是代表节点）
				nodes.put(cur, node);
				parents.put(node, node);
				sizeMap.put(node, 1);
			}
		}

		// 从点cur开始，一直往上找，找到不能再往上的代表点，返回
		public Node<V> findFather(Node<V> cur) {
			// 记录往上找的路径
			Stack<Node<V>> path = new Stack<>();
			while (cur != parents.get(cur)) {
				path.push(cur);
				cur = parents.get(cur);
			}
			// 将往上找的路径上的节点，直接挂在找到的代表节点下（父直接指向代表节点）
			// cur为头节点（代表节点）
			while (!path.isEmpty()) {
				parents.put(path.pop(), cur);
			}
			return cur;
		}

		public boolean isSameSet(V a, V b) {
			if (!nodes.containsKey(a) || !nodes.containsKey(b)) {// 登记了的节点才去找
				return false;
			}
			return findFather(nodes.get(a)) == findFather(nodes.get(b));
		}

		public void union(V a, V b) {
			if (!nodes.containsKey(a) || !nodes.containsKey(b)) {
				return;
			}
			// 找到集合a的代表节点aHead
			Node<V> aHead = findFather(nodes.get(a));
			// 找到集合b的代表节点bHead
			Node<V> bHead = findFather(nodes.get(b));
			if (aHead != bHead) {
			    // 集合a的size
				int aSetSize = sizeMap.get(aHead);
                // 集合b的size
				int bSetSize = sizeMap.get(bHead);
				// 小集合挂在大集合下面：
				// （1）找到大集合的代表点
				Node<V> big = aSetSize >= bSetSize ? aHead : bHead;
				// （2）找到小集合的代表点
				Node<V> small = big == aHead ? bHead : aHead;
				// （3）小挂大
				parents.put(small, big);
				// （4）更新大集合的节点数量
				sizeMap.put(big, aSetSize + bSetSize);
				// （5）移除小集合的size信息
				sizeMap.remove(small);
			}
		}
	}

}
